#include "cuda_fp16.h"
#include "tensor_cuda.h"
#include <algorithm>

template <typename T>
__global__ void fillCudaMemKernel(void *data, const T v, const size_t size) {
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size;
       i += blockDim.x * gridDim.x) {
    reinterpret_cast<T *>(data)[i] = v;
  }
}

template <typename T>
void fillCudaMemAsync(void *data, const T v, const size_t size,
                      cudaStream_t stream) {
  int block = 512;
  int grid = std::min(int(size + block - 1) / block, 1024);
  fillCudaMemKernel<T><<<grid, block, 0, stream>>>(data, v, size);
}

#define INSTANCE(T)                                                            \
  template void fillCudaMemAsync<T>(void *data, const T, const size_t,         \
                                    cudaStream_t);

INSTANCE(float)
INSTANCE(half)
INSTANCE(int8_t)
INSTANCE(int32_t)
INSTANCE(bool)
